from functools import wraps
from fastapi import Request
from core.constants import TimeUnit, CacheConstants
from core.utils import md5
from core.storage import RedisStorage
from core.response import fail
import socket


class LimitType:
    DEFAULT = 1
    IP = 2
    CLUSTER = 3


def RateLimiter(
    key: str = '',
    period: int = 60,
    count: int = 100,
    limitType: int = LimitType.IP,
    timeUnit: int = TimeUnit.SECONDS,
    msg: str = '访问频率过高，请稍后重试'
):
    """
     限流
    :param key:限流key
    :param period:限流时间
    :param count:限流次数
    :param limitType:限流类型
    :param timeUnit:时间单位
    :param msg:提示消息
    :return:
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            req: Request = kwargs.get('req')
            # 获取缓存key
            if limitType == LimitType.IP:
                combineKey = CacheConstants.RATE_LIMIT_KEY + md5(key) + ':' + req.client.host
            elif limitType == LimitType.CLUSTER:
                combineKey = CacheConstants.RATE_LIMIT_KEY + md5(key) + ':' + socket.gethostname() + ":" + req.client.host
            else:
                combineKey = CacheConstants.RATE_LIMIT_KEY + md5(key)

            number = await RedisStorage.get_cache_object(req.app, combineKey)
            if number is None:
                await RedisStorage.set_cache_object(req.app, combineKey, str(count), period * timeUnit)
            elif int(number) <= 0:
                return fail(msg)
            await RedisStorage.decrease(req.app, combineKey)
            return await func(*args, **kwargs)

        return wrapper

    return decorator


def RepeatSubmit(
    interval: int = 1,
    timeUnit: int = TimeUnit.SECONDS,
    msg: str = '提交频率过高，请稍后重试'
):
    """
    自定义注解防止表单重复提交'

    interval:间隔时间(ms)，小于此时间视为重复提交
    timeUnit: 时间单位 s秒
    msg:提示消息 支持国际化 格式为 {code}

    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            req: Request = kwargs.get('req')
            content_type: str = req.headers.get("Content-Type")
            if content_type and ("multipart/form-data" in content_type or 'application/x-www-form-urlencoded' in content_type):
                payload = await req.form()
            else:
                payload = await req.body()

            combineKey = CacheConstants.REPEAT_SUBMIT_KEY + req.client.host + ':' + req.url.path.replace("/","_") + ':' + md5(str(payload, 'utf-8'))
            number = await RedisStorage.get_cache_object(req.app, combineKey)
            if number is None:
                await RedisStorage.set_cache_object(req.app, combineKey, "1", interval * timeUnit)
            else:
                return fail(msg)

            return await func(*args, **kwargs)

        return wrapper

    return decorator


def Sensitive():
    """
        敏感词过滤
    """

    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            req: Request = kwargs.get('req')

            print("---------------过滤敏感词:")
            result = await func(*args, **kwargs)

            return result

        return wrapper

    return decorator
